"""
结果可视化脚本（空壳示例）+

功能：
- 读取 config.yaml 中的 plot 配置（如果可用）
- 使用 utils.plot_style 中的 set_chinese_font 和 set_plot_style
- 画一个简单的示例图并保存到 figures/ 目录

此脚本应作为示例或快速验证使用。
"""
from pathlib import Path
import sys
from rich.console import Console

# 确保项目根目录在 sys.path 中，便于导入 utils 等顶级模块
project_root = Path(__file__).resolve().parents[1]
if str(project_root) not in sys.path:
	sys.path.insert(0, str(project_root))

try:
	import yaml
except Exception:
	yaml = None

# 示例绘图已移除：如需绘图，请在本模块中添加具体绘图函数并调用 utils.plot_style 中的样式函数

from utils.plot_style import set_chinese_font, set_plot_style, force_chinese_font


def load_plot_config(cfg_path: str = None):
	if cfg_path is None:
		cfg_path = Path(__file__).resolve().parents[1] / 'config.yaml'
	cfg_path = Path(cfg_path)
	if not cfg_path.exists() or yaml is None:
		return None
	try:
		with cfg_path.open('r', encoding='utf-8') as f:
			cfg = yaml.safe_load(f)
		return cfg.get('plot') if isinstance(cfg, dict) else None
	except Exception:
		return None


def make_sample_plot(out_dir: str = 'figures', filename: str = 'sample_plot.png', style_cfg: dict = None):
	out_dir = Path(out_dir)
	out_dir.mkdir(parents=True, exist_ok=True)

	# 应用绘图风格与中文字体
	try:
		set_plot_style(style_cfg or {})
	except Exception as e:
		Console().log(f"设置绘图风格失败：{e}")
	try:
		set_chinese_font(style_cfg.get('font_family') if style_cfg else 'Heiti TC')
	except Exception as e:
		Console().log(f"设置中文字体失败：{e}")
	# 正弦波示例已删除；保留函数用于占位或作为示例入口。
	Console().log('示例绘图（正弦波）已移除。如需生成图形，请在本模块中实现具体绘图函数。')
	return None


def plot_border_counties_population_heatmap(
	out_dir: str = 'figures',
	filename: str = 'border_population_heatmap.png',
	year: int = 2020,
	sites_cfg_path: str = 'sites.yaml',
	population_parquet: str = 'data/边境城市七普乡镇.parquet',
	geojson_path: str = 'data/云南乡镇.json',
	cmap: str = 'Reds',
	figsize: tuple = (10, 12),
):
	"""绘制边境县乡镇总人口热力图（choropleth），保存为 PNG。

	步骤：
	- 从 sites.yaml 读取边境县名单（`BORDER_COUNTIES` / `COUNTY_SHORT_NAMES_ALL` 或使用 population 表中的 `border_county` 标记）
	- 读取人口 parquet，按乡镇/乡镇代码筛选目标县的乡镇及其 `总人口`（按 year）
	- 读取 GeoJSON，将人口数据合并到 GeoDataFrame（按乡镇名或乡镇代码匹配）
	- 使用 geopandas/matplotlib 绘制 choropleth 并保存到 out_dir/filename
	"""
	from pathlib import Path
	import yaml
	import pandas as pd

	out_dir = Path(out_dir)
	out_dir.mkdir(parents=True, exist_ok=True)

	console = Console()

	# 应用绘图风格与中文字体（如果提供）
	try:
		if style_cfg:
			set_plot_style(style_cfg)
			set_chinese_font(style_cfg.get('font_family'))
		else:
			# 尝试使用默认样式并设置常用中文字体
			set_plot_style()
			set_chinese_font('Heiti TC')
	except Exception as e:
		console.log(f'应用绘图样式/字体失败：{e}')

	# 读取 sites.yaml（若存在）以获取边境县列表
	border_counties = None
	try:
		if Path(sites_cfg_path).exists():
			with open(sites_cfg_path, 'r', encoding='utf-8') as f:
				cfg = yaml.safe_load(f)
			# 优先使用 explicit 列表
			border_counties = cfg.get('BORDER_COUNTIES') or []
			# 如果未提供完整县名列表，尝试从 mapping 中取 keys
			if not border_counties:
				border_counties = list((cfg.get('COUNTY_SHORT_NAMES_ALL') or {}).keys())
	except Exception as e:
		console.log(f'读取 sites.yaml 失败：{e}')

	# 读取人口数据
	try:
		pop_df = pd.read_parquet(population_parquet)
	except Exception as e:
		console.log(f'无法读取人口数据 {population_parquet}: {e}')
		return None

	# 选择年份
	if '年份' in pop_df.columns:
		pop_df_year = pop_df[pop_df['年份'] == year].copy()
	else:
		pop_df_year = pop_df.copy()

	# 如果 sites.yaml 提供了边境县列表，则筛选；否则使用 population 表里的 border_county==1
	if border_counties:
		pop_df_year = pop_df_year[pop_df_year['县'].isin(border_counties)]
	elif 'border_county' in pop_df_year.columns:
		pop_df_year = pop_df_year[pop_df_year['border_county'] == 1]

	if pop_df_year.empty:
		console.log('未找到匹配的边境县人口数据，绘图终止。')
		return None

	# 准备合并字段：尝试乡镇代码优先，其次乡镇名称
	merge_on_code = None
	if '乡镇代码' in pop_df_year.columns:
		pop_df_year['town_code'] = pop_df_year['乡镇代码'].astype(str)
		merge_on_code = 'town_code'
	elif 'tid' in pop_df_year.columns or 'town_code' in pop_df_year.columns:
		col = 'tid' if 'tid' in pop_df_year.columns else 'town_code'
		pop_df_year['town_code'] = pop_df_year[col].astype(str)
		merge_on_code = 'town_code'

	# 读取 geojson
	try:
		import geopandas as gpd
		gj = gpd.read_file(geojson_path)
	except Exception as e:
		console.log(f'导入 geopandas 或读取 GeoJSON 失败：{e}，请确保已安装 geopandas。')
		return None

	# 标准化 geo 属性：尝试找到乡镇代码或名称列
	geo = gj.copy()
	geo_cols = [c for c in geo.columns]
	# 尝试常见乡镇代码字段
	if '乡镇代码' in geo.columns:
		geo['town_code'] = geo['乡镇代码'].astype(str)
	elif 'tid' in geo.columns:
		geo['town_code'] = geo['tid'].astype(str)
	else:
		# 退回到乡镇/乡列名称匹配
		if '乡' in geo.columns:
			geo['town_name'] = geo['乡'].astype(str)
		elif '乡镇' in geo.columns:
			geo['town_name'] = geo['乡镇'].astype(str)
		elif '名称' in geo.columns:
			geo['town_name'] = geo['名称'].astype(str)

	# 如果可用且 sites.yaml 提供了边境县列表，则优先在 geo 上按县名过滤
	if border_counties:
		county_col = None
		for cand in ['县', 'county', 'COUNTY', '县级']:
			if cand in geo.columns:
				county_col = cand
				break
		if county_col:
			try:
				geo = geo[geo[county_col].isin(border_counties)].copy()
			except Exception:
				# 忽略过滤失败，后续用合并结果筛选
				pass

	# 合并
	if merge_on_code and 'town_code' in geo.columns:
		merged = geo.merge(pop_df_year[['town_code', '总人口']].drop_duplicates('town_code'), on='town_code', how='left')
	else:
		# try name-based merge
		name_field = 'town_name' if 'town_name' in geo.columns else ('乡' if '乡' in geo.columns else None)
		if name_field is None:
			console.log('无法在 GeoJSON 中找到乡镇名称字段用于合并。')
			return None
		merged = geo.merge(pop_df_year[['乡镇', '总人口']].drop_duplicates('乡镇'), left_on=name_field, right_on='乡镇', how='left')

	# 填充为0的乡镇总人口
	merged['总人口'] = merged['总人口'].fillna(0)

	# 如果之前在 geo 上没有成功按县名过滤，则只保留合并后有匹配人口的记录（即边境县乡镇）
	try:
		# 判断是否已经在 geo 上做了县过滤（通过检查 geo 行数与 merged 行数）
		if '县' not in geo.columns and merged['总人口'].sum() > 0:
			merged = merged[merged['总人口'] > 0].copy()
	except Exception:
		pass

	# 绘图
	try:
		ax = merged.plot(column='总人口', cmap=cmap, legend=True, figsize=figsize, edgecolor='black', linewidth=0.2)
		ax.set_axis_off()
		title = f'边境县乡镇总人口热力图（{year}）'
		ax.set_title(title, fontsize=14)
		# 强制设置中文字体到 axes/legend，以避免部分文本仍使用缺省字体
		try:
			force_chinese_font(ax, style_cfg.get('font_family') if style_cfg else 'Heiti TC')
		except Exception:
			pass
		out_path = out_dir / filename
		# 保存 PNG
		fig = ax.get_figure()
		fig.savefig(out_path, dpi=150, bbox_inches='tight')
		console.log(f'已保存热力图到 {out_path}')
		return str(out_path)
	except Exception as e:
		console.log(f'绘图失败：{e}')
		return None


def main():
	# 尝试从项目根读取 config.yaml
	repo_root = Path(__file__).resolve().parents[1]
	cfg_path = repo_root / 'config.yaml'
	style_cfg = load_plot_config(cfg_path) if cfg_path.exists() else None

	figures_dir = repo_root / (style_cfg.get('figures_dir') if style_cfg and style_cfg.get('figures_dir') else 'figures')
	# 默认仅输出占位示例。按需取消注释下面的行以生成边境县人口热力图。
	# make_sample_plot(out_dir=figures_dir, style_cfg=style_cfg)

	# 当脚本通过命令行运行并带有 --heatmap 参数时，生成热力图
	import argparse
	parser = argparse.ArgumentParser(description='plot utilities')
	parser.add_argument('--heatmap', action='store_true', help='生成边境县乡镇人口热力图并保存')
	parser.add_argument('--year', type=int, default=2020, help='年份，默认 2020')
	args = parser.parse_args()
	if args.heatmap:
		plot_border_counties_population_heatmap(out_dir=figures_dir, year=args.year)


if __name__ == '__main__':
	main()

